from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent
from langchain_core.prompts import PromptTemplate
from langchain.tools import tool

from src.module.Ark import ArkModel

## 需要安装 langgraph-checkpoint-redis 后才可以使用
from langgraph.checkpoint.redis import RedisSaver

from langchain_core.output_parsers import StrOutputParser
from langgraph.checkpoint.memory import MemorySaver

class AgentHistory:

    def __init__(self):
        self.model = ArkModel().model
        self.tools = [self.colors]

    @staticmethod
    @tool
    def colors():
        """ 用来获取小皮球的颜色 """
        return '蓝色'

    @property
    def prompt(self):
        return PromptTemplate.from_template("""
            你是一个聊天小助手，用于和用户进行对话。回答用户的问题

            你有以下的工具可以使用，在面对你不知道的问题时，你可以选择你的工具来帮助你。
            请合理的使用你的工具，你的工具如下：{tools}

            用户此次的问题为: {messages}
        """).partial(
            tools= ",".join(item.name for item in self.tools),
        )

    def redis_saver(self):
        with RedisSaver.from_conn_string('redis://127.0.0.1:6379/0') as checkpointer:
            checkpointer.setup()

            agent_execute = create_react_agent(
                self.model,
                self.tools,
                prompt=self.prompt,
                checkpointer=checkpointer
            )

            for chunk in agent_execute.stream(
                {"messages": "我的上个问题是什么?"},
                config={"configurable": {"thread_id": "abc123"}},
                stream_mode="values"
            ):
                chunk["messages"][-1].pretty_print()

    def memory_server(self):
        ## 很慢，走内存都会很慢
        agent_execute = create_react_agent(
            self.model,
            self.tools,
            prompt=self.prompt,
            checkpointer=MemorySaver()
        )

        for chunk in agent_execute.stream(
            {"messages": "我的上个问题是什么?"},
            config={"configurable": {"thread_id": "abc123"}},
            stream_mode="values"
        ):
            chunk["messages"][-1].pretty_print()

    def chain_agent(self):
        model = self.model.bind_tools(self.tools)
        messages = [HumanMessage(content='你的版本是多少？请给我介绍下苹果的生长周期，并回答我小皮球是什么颜色的?')]

        ## langchain 中，bind-tools之后，模型只会调用工具，并不会把调用工具之后的信息进行输出。
        ## 正常的设想是: 用户输出 -> 模型(分析用户的输入) -> 工具(模型调用工具) -> 模型(根据工具返回的内容汇总) -> 输出
        ## 此时：用户输出 -> 模型(分析用户的输入) -> 分析准备调用哪个工具（还未调用工具）
        ai_message = model.invoke(messages)

        ## 手动获取分析的结果
        for item in ai_message.tool_calls:
            select_tool = next((item_tool for item_tool in self.tools if item_tool.name == item['name']), None)
            if not select_tool:
                raise 'un tool call'
            else:
                result = select_tool.invoke(item)
                messages.append(result)

        print(model.invoke(messages))

    def start(self):
        # self.memory_server()
        self.redis_saver()
        # self.chain_agent()